use crate::utils::maybe_deser_owned::VariantMaybeDeserOwned;
use serde::{Deserialize, Serialize};
use std::{
    io::{Error as IoErr, Write},
    marker::PhantomData,
    sync::Arc,
};
use tokio_util::{
    bytes::{Buf, BufMut, Bytes, BytesMut},
    codec::{Decoder, Encoder},
};

macro_rules! unsafe_get_mut {
    ($it: ident, $idx: expr) => {
        unsafe { $it.get_unchecked_mut($idx) }
    };
}

#[cfg(feature = "msgpack")]
pub type DefaultSer = msgpack::DefaultMsgPackSerder;
#[cfg(feature = "msgpack")]
pub type DefaultDeser = msgpack::DefaultMsgPackSerder;
#[cfg(feature = "msgpack")]
pub mod msgpack {
    pub use rmp_serde;

    use super::{FromSliceDeserializer, ToWriterSerializer};
    use rmp_serde::{decode, encode};
    use serde::{Deserialize, Serialize};
    use std::io::Write;

    #[derive(Default, Clone)]
    pub struct DefaultMsgPackSerder;

    impl ToWriterSerializer for DefaultMsgPackSerder {
        type Err = encode::Error;

        fn serialize_to_writer<W: Write, T: Serialize>(
            &self,
            mut w: W,
            t: T,
        ) -> Result<(), Self::Err> {
            encode::write(&mut w, &t)
        }
    }
    impl FromSliceDeserializer for DefaultMsgPackSerder {
        type Err = decode::Error;

        fn deserialize_from_slice<'de, T: Deserialize<'de>>(
            &self,
            slice: &'de [u8],
        ) -> Result<T, Self::Err> {
            decode::from_slice(slice)
        }
    }
}

pub trait FromSliceDeserializer {
    type Err;
    fn deserialize_from_slice<'de, T: Deserialize<'de>>(
        &self,
        slice: &'de [u8],
    ) -> Result<T, Self::Err>;
}

impl<D: FromSliceDeserializer> FromSliceDeserializer for Arc<D> {
    type Err = D::Err;
    fn deserialize_from_slice<'de, T: Deserialize<'de>>(
        &self,
        slice: &'de [u8],
    ) -> Result<T, Self::Err> {
        D::deserialize_from_slice(self, slice)
    }
}

pub trait ToWriterSerializer {
    type Err;
    fn serialize_to_writer<W: Write, T: Serialize>(&self, w: W, t: T) -> Result<(), Self::Err>;
}
impl<S: ToWriterSerializer> ToWriterSerializer for Arc<S> {
    type Err = S::Err;
    fn serialize_to_writer<W: Write, T: Serialize>(&self, w: W, t: T) -> Result<(), Self::Err> {
        S::serialize_to_writer(self, w, t)
    }
}

pub(crate) trait FixedSizedPrefix {
    type Err;
    const SIZE: usize;
    /// Write the data this structure contains to the given byte
    /// buffer.
    fn write_to(self, buf: &mut [u8]) -> Result<(), Self::Err>;
    /// Construct self from the given byte buffer.
    ///
    /// # Return
    ///
    /// - Whether deserialization of payload of should be delegated
    ///    to code external to the codec.
    /// - Constructed prefix structure.
    fn read_from(buf: &[u8]) -> Result<(bool, Self), Self::Err>
    where
        Self: Sized;
}

pub(crate) struct WithFixedSizedPrefix<Ser, Deser, EncPre, DecPre, I, O> {
    enc_ser: Ser,
    dec_deser: Deser,
    dec_size: Option<u32>,
    marker: PhantomData<(EncPre, DecPre, I, O)>,
}

impl<Ser, Deser, EncPre, DecPre, I, O> WithFixedSizedPrefix<Ser, Deser, EncPre, DecPre, I, O> {
    pub(crate) fn new(enc_ser: Ser, dec_deser: Deser) -> Self {
        Self {
            enc_ser,
            dec_deser,
            dec_size: None,
            marker: PhantomData,
        }
    }
}

pub(crate) enum IoOrDecodeErr<PE, P, DE> {
    Io(IoErr),
    ShouldDelegate,
    PrefixRead(PE),
    PayloadDeser(P, DE),
}
impl<PE, P, DE> From<IoErr> for IoOrDecodeErr<PE, P, DE> {
    fn from(value: IoErr) -> Self {
        Self::Io(value)
    }
}
pub(crate) enum IoOrEncodeErr<PE, SE> {
    Io(IoErr),
    PrefixWrite(PE),
    PayloadSer(SE),
}
impl<PE, SE> IoOrEncodeErr<PE, SE> {
    pub(crate) fn assert_io_err(self) -> IoErr {
        if let IoOrEncodeErr::Io(e) = self {
            e
        } else {
            panic!("Assert fail: IoOrEncodeErr isn't IO Error.")
        }
    }
}
impl<PE, SE> From<IoErr> for IoOrEncodeErr<PE, SE> {
    fn from(value: IoErr) -> Self {
        Self::Io(value)
    }
}

/// Like `Option<T>`, but untagged when being serialized, so when
/// it's `None`, it will be serialized to nothing.
#[derive(Serialize)]
#[serde(untagged)]
pub(crate) enum MaybeEmpty<T> {
    Filled(T),
    Empty,
}

impl<Ser, Deser, EncPre, DecPre, I, O> Encoder<(EncPre, MaybeEmpty<I>)>
    for WithFixedSizedPrefix<Ser, Deser, EncPre, DecPre, I, O>
where
    Ser: ToWriterSerializer,
    EncPre: FixedSizedPrefix,
    I: Serialize,
{
    type Error = IoOrEncodeErr<EncPre::Err, Ser::Err>;

    fn encode(
        &mut self,
        (pre, item): (EncPre, MaybeEmpty<I>),
        dst: &mut BytesMut,
    ) -> Result<(), Self::Error> {
        let orig_len = dst.len();
        dst.reserve(4 + EncPre::SIZE);
        // SAFETY: we never read these bytes before their initialization
        unsafe { dst.advance_mut(4 + EncPre::SIZE) };
        let w = dst.writer();
        if let Err(e) = self.enc_ser.serialize_to_writer(w, item) {
            // SAFETY: we simply retreat back
            unsafe { dst.set_len(orig_len) };
            return Err(IoOrEncodeErr::PayloadSer(e));
        }
        let size = (dst.len() - orig_len - 4) as u32;
        // SAFETY: we have reserved at least (4 + EncPre::SIZE) bytes after `orig_len`.
        let size_field = unsafe_get_mut!(dst, orig_len..orig_len + 4);
        size_field.copy_from_slice(&(size.to_le_bytes()));
        let other_prefix = unsafe_get_mut!(dst, orig_len + 4..orig_len + 4 + EncPre::SIZE);
        pre.write_to(other_prefix)
            .map_err(IoOrEncodeErr::PrefixWrite)?;

        Ok(())
    }
}

pub(crate) enum Decoded<T> {
    Empty,
    Source(Bytes),
    Deserialzed(T),
}

impl<Ser, Deser, EncPre, DecPre, I, O> Decoder
    for WithFixedSizedPrefix<Ser, Deser, EncPre, DecPre, I, O>
where
    Deser: FromSliceDeserializer,
    DecPre: FixedSizedPrefix,
    O: Deserialize<'static> + VariantMaybeDeserOwned,
{
    type Item = (DecPre, Decoded<O>);
    type Error = IoOrDecodeErr<DecPre::Err, DecPre, Deser::Err>;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        loop {
            if let Some(size) = self.dec_size {
                let size = size as usize;
                let data: &[u8] = &*src;
                if data.len() < size {
                    // note that we remain in the 'waiting for enough data to
                    // decode after the size field' state when we return here,
                    // so the size field should not be reset until the check
                    // above passes.
                    return Ok(None);
                }
                assert!(
                    size >= DecPre::SIZE,
                    "frame size ({size}) < fixed prefix length ({}), no recovery method is implemented, exiting.",
                    DecPre::SIZE
                );
                self.dec_size = None;
                match DecPre::read_from(&data[..DecPre::SIZE]) {
                    Err(e) => {
                        // Skip this frame
                        src.advance(size);
                        return Err(IoOrDecodeErr::PrefixRead(e));
                    }
                    Ok((delegate_deser, pre)) => {
                        if size == DecPre::SIZE {
                            src.advance(size);
                            return Ok(Some((pre, Decoded::Empty)));
                        }
                        if delegate_deser {
                            let mut out = src.split_to(size).freeze();
                            out.advance(DecPre::SIZE);
                            return Ok(Some((pre, Decoded::Source(out))));
                        } else {
                            let payload = &data[DecPre::SIZE..size];

                            // SAFETY:
                            // safety here depends on:
                            // 1. `deserialize_from_slice<'a, 'de>(&'a self, src: &'de [u8])`
                            //    does not leak the reference to `src` into outside world.
                            // 2. `variant_is_deser_owned` implementation does not leak
                            //    reference to `self` into outside world.
                            // 3. `variant_is_deser_owned` contains correct implementation,
                            //    i.e. when it returns `true`, the value desered indeed does
                            //    not contain any reference to the outside world.
                            let payload_as_static: &'static [u8] =
                                unsafe { std::mem::transmute(payload) };
                            let res = self
                                .dec_deser
                                .deserialize_from_slice::<O>(payload_as_static);
                            src.advance(size);
                            return match res {
                                Err(e) => Err(IoOrDecodeErr::PayloadDeser(pre, e)),
                                Ok(r) if r.variant_is_deser_owned() => {
                                    Ok(Some((pre, Decoded::Deserialzed(r))))
                                }
                                _ => Err(IoOrDecodeErr::ShouldDelegate),
                            };
                        }
                    }
                }
            } else if src.len() >= 4 {
                let size = src.get_u32_le();
                src.reserve(size as usize);
                self.dec_size = Some(size);
            } else {
                return Ok(None);
            }
        }
    }
}
